package cn.jbolt.ai.klb.service.impl;


import cn.hutool.core.date.DateTime;
import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.core.fileSplit.FileSplitService;
import cn.jbolt.ai.core.vdb.VdbCollectionService;
import cn.jbolt.ai.klb.DocDataType;
import cn.jbolt.ai.klb.DocLinkType;
import cn.jbolt.ai.klb.entity.AiDoc;
import cn.jbolt.ai.klb.service.AiDocService;
import cn.jbolt.util.JSONUtil;
import cn.jbolt.util.Pager;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.UpdateWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import cn.jbolt.ai.klb.FileState;
import cn.jbolt.ai.klb.FileType;
import cn.jbolt.ai.klb.dao.AiFileMapper;
import cn.jbolt.ai.klb.entity.AiFile;
import cn.jbolt.ai.klb.entity.AiKnowledgeBase;
import cn.jbolt.ai.klb.service.AiFileService;
import cn.jbolt.ai.klb.service.AiKnowledgeBaseService;
import cn.jbolt.ai.klb.vo.AiFileListAddVO;
import cn.jbolt.util.Result;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.jboltai.capability.vdb.search.VDBConditions;
import com.jboltai.resource.ai.AIModel;
import org.apache.poi.util.StringUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.List;
import java.util.stream.Collectors;


@Service
public class AiFileServiceImpl extends ServiceImpl<AiFileMapper, AiFile> implements AiFileService {

    @Autowired
    private AiKnowledgeBaseService aiKnowledgeBaseService;
    @Autowired
    private FileSplitService fileSplitService;

    @Autowired
    private AiDocService aiDocService;
    @Autowired
    private VdbCollectionService vdbCollectionService;

    @Autowired
    @Qualifier("threadPoolTaskExecutor")
    private ThreadPoolTaskExecutor threadPool;

    private static final Logger logger = LoggerFactory.getLogger(AiFileServiceImpl.class);

    @Override
    public Result<AiFile> getFileById(Long id) {
        AiFile aiFile = this.getById(id);
        if (aiFile != null) {
            aiFile = getKnowledgeBaseName(aiFile);
            return Result.success(aiFile);
        } else {
            return Result.error("文件不存在");
        }
    }

    @Override
    public Result listByPage(Integer pageNumber, Integer pageSize, String keywords, Long knowledgeBaseId, String trainType,
                             Integer status) {

        QueryWrapper<AiFile> queryWrapper = new QueryWrapper<>();
        if (StrUtil.isNotBlank(keywords)) queryWrapper.like("file_name", keywords);
        if (knowledgeBaseId != null) queryWrapper.eq("knowledge_base_id", knowledgeBaseId);
        if (StrUtil.isNotBlank(trainType)) queryWrapper.eq("doc_type", trainType);
        if (status != null) queryWrapper.eq("file_status", status);
        queryWrapper.orderByDesc("id");
        Page<AiFile> pager = this.page(new Page<>(pageNumber, pageSize), queryWrapper);
        List<AiFile> fileList = pager.getRecords();
        fileList.stream().map(aiFile -> {
            aiFile = getKnowledgeBaseName(aiFile);
            return aiFile;
        }).collect(Collectors.toList());
        pager.setRecords(fileList);
        return Result.success(Pager.from(pager));
    }


    @Override
    public Result batchAdd(AiFileListAddVO vo) {
        Result result = aiKnowledgeBaseService.checkKnowledgeBase(Long.valueOf(vo.getKnowledgeBaseId()));
        if (result.isError()) return result;
        List<AiFile> fileList = vo.getList();
        if (fileList == null || fileList.size() == 0) return Result.error("请选择文件");
        //只有QA文档才需要判断ai模型
        AIModel aiModel = null;
        String defaultModel = null;
        if (FileType.QA.getCode().equals(vo.getDocType())) {
            if (StrUtil.isNotBlank(vo.getModel())) {
                // 使用 indexOf() 找到第一个冒号的位置
                int index = vo.getModel().indexOf(':');
                // 使用 substring() 方法来获取两个部分
                String modelType = vo.getModel().substring(0, index);
                if (StrUtil.isNotBlank(modelType)) {
                    aiModel = AIModel.get(modelType);
                }
                //取出defaultModel
                defaultModel = vo.getModel().substring(index + 1);

            }
        }
        //在循环外部构建规则
        ObjectNode objectNode = JSONUtil.createObjectNode();
        if (vo.getChunkSize() != null && vo.getChunkSize() != 0) {
            objectNode.put("chunkSize", vo.getChunkSize());
        }
        if (vo.getChunkOverlap() != null && vo.getChunkOverlap() != 0) {
            objectNode.put("chunkOverlap", vo.getChunkOverlap());
        }

        if (aiModel != null) {
            objectNode.put("model", vo.getModel());
        }
        if (vo.getExtractResources() != null) {
            objectNode.put("extractResources", vo.getExtractResources());
        }
        if (vo.getResourcesPosition() != null) {
            objectNode.put("resourcesPosition", vo.getResourcesPosition());
        }
        if (vo.getResourceSavePath() != null) {
            objectNode.put("resourceSavePath ", vo.getResourceSavePath());
        }
        objectNode.put("ocrImg", vo.getOcrImg());
        String json = JSONUtil.toJson(objectNode);
        for (AiFile aiFile : fileList) {
            aiFile.setChunkSize(vo.getChunkSize());
            aiFile.setChunkOverlap(vo.getChunkOverlap());
            aiFile.setAiModel(aiModel);
            aiFile.setDefaultModel(defaultModel);
            aiFile.setKnowledgeBaseId(vo.getKnowledgeBaseId());
            aiFile.setFileStatus(FileState.NO_TRAIN.getCode());
            aiFile.setSplitRule(json);
        }

        boolean b = this.saveBatch(fileList);
        if (b) {
            //文件训练+向量化
            for (AiFile aiFile : fileList) {
                threadPool.submit(() -> trainFile(aiFile));
            }
            return Result.success();
        } else {
            return Result.error("添加失败");
        }
    }

    /**
     * 训练文件
     *
     * @return
     */
    @Transactional
    public void trainFile(AiFile file) {
        file.setFileStatus(FileState.TRAINING.getCode()); //设置状态为训练中
        this.updateById(file);
        Result fileResult = fileSplitService.split(file);//训练文件中
        if (fileResult.isError()) {
            file.setFileStatus(FileState.TRAIN_ERROR.getCode());//训练失败的状态
            this.updateById(file);
        } else {
            //有一种情况,是大模型拆分的时候,没有拆出DOC,直接就把file的状态设置成失败了
            if (this.getFileById(file.getId()).getState().equals(FileState.TRAINING.getCode())) {
                file.setTrainingTime(DateTime.now());
                file.setFileStatus(FileState.TRAIN_SUCCESS.getCode());//训练成功的状态
                this.updateById(file);
            }
        }


    }


    @Override
    public void updateState(Long fileId, Integer state) {
        UpdateWrapper<AiFile> updateWrapper = new UpdateWrapper<>();
        updateWrapper.eq("id", fileId);
        updateWrapper.set("file_status", state);
        if (state.equals(FileState.EMBEDDING_ERROR.getCode())) {//如果是失败状态的话
            updateWrapper.set("embed_time", null);
        } else if (state.equals(FileState.EMBEDDING_SUCCESS.getCode())) {//如果是成功状态的话
            updateWrapper.set("embed_time", DateTime.now());
        }
        this.update(updateWrapper);
    }


    @Override
    public Result update(AiFile file) {
        if (updateById(file)) {
            return Result.success();
        } else {
            return Result.error("更新失败");
        }
    }

    @Override
    @Transactional
    public Result delete(Long id) {
        Result<AiFile> fileResult = getFileById(id);
        if (fileResult.isError()) return fileResult;
        AiFile file = fileResult.getData();
        if (removeById(id)) {
            //移除doc
            Result result = aiDocService.delDocByFileId(id);
            if (result.isError()) return result;
            //移除相应的vdb
            Result knowledgeBaseResult = aiKnowledgeBaseService.getKnowledgeBaseById(file.getKnowledgeBaseId());
            if (knowledgeBaseResult.isError()) return knowledgeBaseResult;
            AiKnowledgeBase knowledgeBase = (AiKnowledgeBase) knowledgeBaseResult.getData();
            VDBConditions conditions = new VDBConditions();
            conditions.eq("file_id", id.toString());
            Result removeVDBResult = vdbCollectionService.removeDataByConditions(knowledgeBase.getVdbResourceId().toString(),
                    "KDB_" + knowledgeBase.getId(), conditions);
            if (removeVDBResult.isError()) return Result.error("移除vdb失败");
        } else {
            return Result.error("删除失败");
        }
        return Result.success();
    }


    public AiFile getKnowledgeBaseName(AiFile aiFile) {
        if (aiFile.getKnowledgeBaseId() != null) {
            Result result = aiKnowledgeBaseService.getKnowledgeBaseById(aiFile.getKnowledgeBaseId());
            if (result.isSuccess()) {
                AiKnowledgeBase aiKnowledgeBase = (AiKnowledgeBase) result.getData();
                aiFile.setAiKnowledgeBaseName(aiKnowledgeBase.getName());
            }
        }
        return aiFile;
    }

    @Override
    public Result getTypeList() {
        return Result.success(FileState.toList());

    }

    @Override
    public Result getTrainList() {
        return Result.success(FileType.toList());
    }


    /**
     * 根据知识库id移除file
     *
     * @return
     */
    @Override
    public Result delFileByKlbId(Long klbId) {
        QueryWrapper<AiFile> fileQuery = new QueryWrapper<AiFile>().eq("knowledge_base_id", klbId);
        List<AiFile> fileList = this.list(fileQuery);
        if (fileList != null && fileList.size() > 0) {
            boolean remove = this.remove(fileQuery);
            if (remove) {
                return Result.success();
            }
            return Result.error("删除失败");
        }
        return Result.success();
    }

    /**
     * 重新训练文件
     *
     * @param id
     * @return
     */
    @Override
    public Result reset(Long id) {
        AiFile file = getById(id);
        if (file == null) {
            return Result.error("文件不存在");
        }
        Result result = aiKnowledgeBaseService.checkKnowledgeBase(file.getKnowledgeBaseId());
        if (result.isError()) return result;
        //把规则重新赋值到file上
        String splitRule = file.getSplitRule();
        if (StringUtil.isNotBlank(splitRule)) {
            ObjectNode jsonObject = JSONUtil.toJsonObject(splitRule);
            if (jsonObject != null) {
                JsonNode chunkSize = jsonObject.get("chunkSize");
                if (chunkSize != null) {
                    file.setChunkSize(chunkSize.asInt());
                }
                JsonNode chunkOverlap = jsonObject.get("chunkOverlap");
                if (chunkOverlap != null) {
                    file.setChunkOverlap(chunkOverlap.asInt());
                }
                JsonNode model = jsonObject.get("model");
                if (model != null) {
                    file.setAiModel(AIModel.get(model.asText()));
                }
                JsonNode extractResources = jsonObject.get("extractResources");
                if (extractResources != null) {
                    file.setExtractResources(extractResources.asBoolean());
                }
                JsonNode resourcesPosition = jsonObject.get("resourcesPosition");
                if (resourcesPosition != null) {
                    file.setResourcesPosition(resourcesPosition.asText());
                }
                file.setOcrImg(jsonObject.has("ocrImg") ? jsonObject.get("ocrImg").asBoolean(): false);
                //这个地方是为了兼容旧数据,之前的时候字段名是path,现在字段名是resourceSavePath,所以这里兼容一下
                JsonNode path = jsonObject.get("path");
                if (path != null) {
                    file.setResourceSavePath(path.asText());
                } else {
                    JsonNode resourceSavePath = jsonObject.get("resourceSavePath");
                    if (resourceSavePath != null) {
                        file.setResourceSavePath(resourceSavePath.asText());
                    }
                }

                //文件训练+向量化
                threadPool.submit(() -> trainFile(file));
                return Result.success();
            }

            return Result.error("该文件无法重新训练");

        } else {
            return Result.error("该文件无法重新训练");
        }
    }

    @Override
    public Result batchAddQuestionByFile(String modelSn, AiFile file) {
        if (StrUtil.isEmpty(modelSn)) {
            return Result.error("请选择大模型");
        }
        if (file == null) {
            return Result.error("文件不存在");
        }
        List<AiDoc> aiDocList = processFile(file);
        if (aiDocList == null || aiDocList.size() == 0) {
            return Result.error("未查询到相关的DOC");
        }

        for (AiDoc doc : aiDocList) {

            AiDoc newDoc = createNewAiDoc(doc);

            submitTask(modelSn, file.getDocType(), newDoc);
        }

        return Result.success();
    }

    private AiDoc createNewAiDoc(AiDoc doc) {
        AiDoc newDoc = new AiDoc();
        newDoc.setContent(doc.getContent());
        newDoc.setAnswer(doc.getDocType().equals(FileType.QA.getCode()) ? doc.getAnswer() : null);
        newDoc.setKnowledgeBaseId(doc.getKnowledgeBaseId());
        newDoc.setFileId(doc.getFileId());
        newDoc.setDocId(doc.getId());
        newDoc.setDataType(DocDataType.CUSTOM_TEXT.getCode());
        newDoc.setDocStatus(FileState.TRAIN_SUCCESS.getCode());
        newDoc.setDocType(FileType.QA.getCode());
        newDoc.setTrainTime(DateTime.now());
        newDoc.setLinkId(doc.getDocType().equals(FileType.QA.getCode()) ? doc.getLinkId() : doc.getId());
        newDoc.setLinkType(doc.getDocType().equals(FileType.QA.getCode()) ? doc.getLinkType() : DocLinkType.LINK_TYPE_DOC.getCode());
        return newDoc;
    }

    /**
     * 提交任务到线程池进行异步处理
     *
     * @param modelSn
     * @param
     * @param newDoc
     */
    private void submitTask(String modelSn, String docType, AiDoc newDoc) {
        try {
            if (docType.equals(FileType.NORMAL.getCode())) {
                threadPool.submit(() -> aiDocService.generateQuestionByModel(modelSn, newDoc.getContent(), newDoc));
            } else {
                threadPool.submit(() -> aiDocService.generateQuestionByModel(modelSn, newDoc.getAnswer(), newDoc));
            }
        } catch (Exception e) {
            logger.error("提交任务到线程池时发生异常", e);
        }
    }

    /**
     * 处理文件并获取相应的文档列表
     *
     * @param
     * @return 文档列表，如果发生错误或异常则返回null
     */
    private List<AiDoc> processFile(AiFile file) {

        if (file.getDocType().equals(FileType.NORMAL.getCode())) { //非QA
            Result<List<AiDoc>> result = aiDocService.getNormalDocListByFileId(file.getId());
            if (result.isError()) {
                logger.error("获取非QA文档列表失败: {}", result.getMsg());
                return null;
            }
            return result.getData();
        } else {
            Result<List<AiDoc>> result = aiDocService.getQADocListByFileId(file.getId());
            if (result.isError()) {
                logger.error("获取QA文档列表失败: {}", result.getMsg());
                return null;
            }
            return result.getData();
        }

    }
}
